# LSTM预测手写数字
import torch
import torch.nn as nn   # 神经网络模块，用于定义各种层（如卷积层、全连接层）和构建神经网络模型。
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST              # 库的一部分，专门用于加载流行的计算机视觉数据集，这里是 MNIST 手写数字数据集。
from torchvision.transforms import ToTensor, Normalize, Compose
from torch.optim import SGD
import torch.nn.init as init # 虽然LSTM层通常有默认初始化，但我们依然可以了解如何手动初始化
from PIL import Image # 导入PIL库用于图像处理
import os
import numpy as np
import matplotlib.pyplot as plt

class Net(nn.Module):       # 定义神经网络
    def __init__(self, input_size=28, hidden_size=128, num_layers=1):
        super(Net, self).__init__()
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, batch_first=True)
        # LSTM层的定义
        # input_size: 每个时间步的特征数，对于28x28的MNIST，我们把每一行看作一个时间步的输入，所以input_size=28
        # hidden_size: 隐藏状态的特征数，这个可以自己定义，类似于CNN的输出通道数
        # num_layers: LSTM层的层数，多层LSTM可以将前一层的输出作为后一层的输入
        # batch_first=True: 这是非常重要的参数。如果设置为True，则输入和输出张量的第一个维度是batch_size。
        #                   即输入形状为 (batch, seq_len, input_size)。
        #                   如果为False (默认)，则输入形状为 (seq_len, batch, input_size)。
        #                   对于图像数据，通常我们更习惯于batch_first=True。
        self.fc = nn.Linear(128, 10) # 全连接层
    def forward(self, x):
        x = x.squeeze(1)  # 去掉通道维度，形状变为 [batch_size, seq_len, input_size]
        _, (h_n, _) = self.lstm(x)  # 通过LSTM层，h_n是最后一个时间步的隐藏状态
        x = h_n[-1]  # 取最后一层的隐藏状态进行分类
        x = self.fc(x)  # 全连接层输出结果
        return x

# 准备数据集 (与CNN部分完全相同，因为数据本身不变)
def prepare_data(path):
    mean = (0.1307)
    std = (0.3081)
    # 定义数据转换：先转换为张量，然后标准化
    trans = Compose([ToTensor(), Normalize((mean,), (std,))])
    # 加载训练数据集，如果不存在则下载
    train = MNIST(path, train=True, download=True, transform=trans)
    # 加载测试数据集，如果不存在则下载
    test = MNIST(path, train=False, download=True, transform=trans)
    # 为训练集创建 DataLoader，batch_size=64，打乱数据
    train_dl = DataLoader(train, batch_size=64, shuffle=True)
    # 为测试集创建 DataLoader，batch_size=1024，不打乱数据（测试时通常不需要）
    test_dl = DataLoader(test, batch_size=1024, shuffle=False)
    return train_dl, test_dl, mean, std

# 训练模型方法 (与CNN部分几乎相同，只是模型是LSTM)
def train_model(train_dl, model):
    crit = nn.CrossEntropyLoss()
    optimizer = SGD(model.parameters(), lr=0.01)
    num_epochs = 5 # 训练轮次
    losses = []
    # 遍历训练数据集中的每个 mini batch
    for epoch in range(num_epochs):
        model.train() # 设置为训练模式
        for i, (inputs, targets) in enumerate(train_dl):
            # 前向传播
            outputs = model(inputs)
            # 计算损失
            loss = crit(outputs, targets)
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
            # 记录损失值
            losses.append(loss.item())
            if i % 100 == 0:
                print(f"Epoch {epoch + 1}/{num_epochs}, Batch {i}/{len(train_dl)}, Loss: {loss.item():.4f}")
    return losses  # 返回所有损失值


# 评估模型方法 (与CNN部分完全相同)
def evaluate_model(model, test_dl):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, targets in test_dl:
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += targets.size(0)
            correct += (predicted == targets).sum().item()
    accuracy = 100 * correct / total
    return accuracy

# 预测单张图片方法 (与CNN部分相同，但要确保图片处理和模型输入匹配LSTM)
# 关键在于 predict_single_image 中的 `img_tensor.unsqueeze(0)` 和 `img_tensor.squeeze(1)` 确保维度正确
def predict_single_image(model, image_path, mean, std):
    model.eval()  # 将模型设置为评估模式

    if not os.path.exists(image_path):
        print(f"错误：文件 '{image_path}' 不存在。")
        return None

    try:
        # 加载图片
        image = Image.open(image_path)

        # 转换为灰度图 (如果不是灰度图)
        if image.mode != 'L':
            print("图片已转换为灰度图。")
            image = image.convert('L')  # 'L'代表灰度图

        # 调整大小到28x28 (MNIST图片的标准尺寸)
        image = image.resize((28, 28))

        # 将PIL图像转换为NumPy数组，用于计算平均像素值
        img_np = np.array(image)

        # 计算平均像素值（0-255范围）
        mean_pixel_value = img_np.mean()
        print(f"图片平均像素值: {mean_pixel_value:.2f}")

        # 因此，如果平均像素值远高于128（例如，超过150），则认为是白底。
        inversion_threshold = 150  # 可以根据实际测试进行调整

        transform_list = [ToTensor()]
        if mean_pixel_value > inversion_threshold:
            print("检测到白底黑字，进行颜色反转。")
            # 在ToTensor之后反转，因为ToTensor将0-255映射到0-1
            transform_list.append(lambda x: 1.0 - x)
        else:
            print("检测到黑底白字或其它，不进行颜色反转。")

        transform_list.append(Normalize(mean, std))

        final_transform = Compose(transform_list)
        img_tensor = final_transform(image)

        # 添加批次维度 (模型期望的输入是 [batch_size, channels, height, width])
        # 对于LSTM，后续会在forward中squeeze掉通道维度
        img_tensor = img_tensor.unsqueeze(0)  # 从 [C, H, W] 变为 [1, C, H, W]

        # 进行预测
        with torch.no_grad():
            output = model(img_tensor) # LSTM模型内部会处理维度转换
            probabilities = F.softmax(output, dim=1)  # 转换为概率分布
            predicted_class = torch.argmax(probabilities, dim=1).item()
            confidence = probabilities[0, predicted_class].item()

        print(f"预测结果：{predicted_class}")
        print(f"置信度：{confidence:.4f}")
        return predicted_class

    except Exception as e:
        print(f"处理图片时发生错误：{e}")
        return None

def plot_losses(losses):
    plt.plot(losses)
    plt.title('Training Loss Curve')
    plt.xlabel('Batch')
    plt.ylabel('Loss')
    plt.show()

if __name__ == "__main__":
    data_path = './data'
    model_save_path = 'mnist_lstm_model.pth' # 修改保存的模型文件名
    train_losses = []  # 用于存储训练损失

    # 准备数据，并获取用于自定义图片预测的mean和std
    train_dl, test_dl, mnist_mean, mnist_std = prepare_data(data_path)

    # 实例化LSTM模型
    # input_size=28 (每一行28像素), hidden_size=128 (隐藏层大小), num_layers=2 (2层LSTM)
    model = Net(input_size=28, hidden_size=128, num_layers=2)

    # 如果模型文件存在，尝试加载
    if os.path.exists(model_save_path):
        try:
            model.load_state_dict(torch.load(model_save_path))
            print(f"已加载保存的模型：{model_save_path}")
        except Exception as e:
            print(f"加载模型失败：{e}。将重新初始化模型。")
    else:
        print("未找到模型文件，将从头开始训练。")

    while True:
        # 选择操作：
        print("\n请选择操作：")
        print("1. 训练模型")
        print("2. 预测自定义图片 (图片名需为 temp.png 或 temp.jpg 等)")
        print("3. 展示训练损失曲线")  # 新增选项
        print("0. 退出")
        choice = input("请输入你的选择 (1/2/3/0): ")

        if choice == '1':
            train_losses = train_model(train_dl, model)  # 训练模型并获取损失值
            print("训练完成！")
            print("开始评估模型在测试集上的性能...")
            test_accuracy = evaluate_model(model, test_dl)
            print(f"模型在测试集上的准确率: {test_accuracy:.2f}%")
            # 训练结束后保存模型
            torch.save(model.state_dict(), model_save_path)
            print(f"模型已保存为 {model_save_path}")
        elif choice == '2':
            if not os.path.exists(model_save_path) and not any(p.grad is not None for p in model.parameters()):
                print("请先训练模型或确保模型文件存在。")
                continue

            image_name = input("请输入图片文件名 (例如 temp.png 或 temp.jpg，图片应在当前目录下): ")
            predicted_digit = predict_single_image(model, image_name, mnist_mean, mnist_std)
            if predicted_digit is not None:
                print(f"模型预测该图片中的数字是：{predicted_digit}")
        elif choice == '3':
            if not train_losses:
                print("没有训练损失数据，请先进行训练。")
            else:
                plot_losses(train_losses)  # 调用绘图方法展示损失曲线
        elif choice == '0':
            print("程序退出。")
            break
        else:
            print("无效的选择，请重新输入。")

